package com.tuniu.agents.advisor;

import com.tuniu.agents.util.JsonUtil;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.api.AdvisedRequest;
import org.springframework.ai.chat.client.advisor.api.AdvisedResponse;
import org.springframework.ai.chat.client.advisor.api.CallAroundAdvisorChain;
import org.springframework.ai.chat.client.advisor.api.StreamAroundAdvisorChain;
import org.springframework.ai.chat.model.MessageAggregator;
import reactor.core.publisher.Flux;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Agent global environment variable storage aspect
 */
public class ChatExtendMemoryAdvisor extends AbstractChatMemoryAdvisor<ChatEnvironmentMemory> {

    public static final String CHAT_MEMORY_EXTEND_MSG = "chat_memory_extend_msg";

    public static final String CHAT_MODEL_LIST = "chat_model_list";

    public ChatExtendMemoryAdvisor(ChatEnvironmentMemory chatEnvironmentMemory) {
        super(chatEnvironmentMemory);
    }

    public ChatExtendMemoryAdvisor(ChatEnvironmentMemory chatEnvironmentMemory, String defaultConversationId, int chatHistoryWindowSize,
                                   int order) {
        super(chatEnvironmentMemory, defaultConversationId, chatHistoryWindowSize, true, order);
    }

    public static ChatExtendMemoryAdvisor.Builder builder(ChatEnvironmentMemory chatEnvironmentMemory) {
        return new ChatExtendMemoryAdvisor.Builder(chatEnvironmentMemory);
    }

    @Override
    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
        advisedRequest = before(advisedRequest);
        AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);
        return advisedResponse;
    }

    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
        advisedRequest = before(advisedRequest);
        Flux<AdvisedResponse> advisedResponses = chain.nextAroundStream(advisedRequest);
        return new MessageAggregator().aggregateAdvisedResponse(advisedResponses, this::observeAfter);
    }

    private AdvisedRequest before(AdvisedRequest request) {

        String conversationId = this.doGetConversationId(request.adviseContext());

        // Bring global information into tool
        Map<String,Object> extendMsg = doGetChatExtendMsg(request.adviseContext());
        if (this.getChatMemoryStore().exist(conversationId)) {
            // Use existing data for additional information, do not overwrite
            extendMsg = this.getChatMemoryStore().get(conversationId);
            // Record model information into global variable
            Object modelObject = this.getChatMemoryStore().get(conversationId, CHAT_MODEL_LIST);
            if(modelObject!=null){
                List<String> modelList = JsonUtil.toList(modelObject.toString(),String.class);
                String modelName = JsonUtil.toMap(JsonUtil.toJson(request.chatOptions())).get("model").toString();
                if(!modelList.contains(modelName)){
                    modelList.add(modelName);
                }
                this.getChatMemoryStore().add(conversationId, Map.of(CHAT_MODEL_LIST,JsonUtil.toJson(modelList)));
            }else{
                List<String> modelList = Arrays.asList(JsonUtil.toMap(JsonUtil.toJson(request.chatOptions())).get("model").toString());
                this.getChatMemoryStore().add(conversationId, Map.of(CHAT_MODEL_LIST,JsonUtil.toJson(modelList)));
            }
        } else {
            if(extendMsg!=null) {
                // Store additional information
                this.getChatMemoryStore().add(conversationId, extendMsg);
            }else {
                extendMsg = new HashMap<>();
            }
            // Record model information into global variable
            List<String> modelList = Arrays.asList(JsonUtil.toMap(JsonUtil.toJson(request.chatOptions())).get("model").toString());
            this.getChatMemoryStore().add(conversationId, Map.of(CHAT_MODEL_LIST,JsonUtil.toJson(modelList)));
        }
        // Add the toolContext that comes with the request itself
        if(request.toolContext()!=null){
            extendMsg.putAll(request.toolContext());
            extendMsg.put(AgentChatMemoryAdvisor.CHAT_MEMORY_AGENT_ID_KEY,doGetAgentId(request.adviseContext()));
            AdvisedRequest advisedRequest = AdvisedRequest.from(request).toolContext(extendMsg).build();
            return advisedRequest;
        }
        return request;
    }

    private void observeAfter(AdvisedResponse advisedResponse) {
    }

    protected Map<String,Object> doGetChatExtendMsg(Map<String, Object> context){
        if(context.containsKey(CHAT_MEMORY_EXTEND_MSG)){
            String content = JsonUtil.toJson(context.get(CHAT_MEMORY_EXTEND_MSG));
            return JsonUtil.toMap(content);
        }else {
            return null;
        }
    }

    protected String doGetAgentId(Map<String, Object> context) {
        return context.containsKey(AgentChatMemoryAdvisor.CHAT_MEMORY_AGENT_ID_KEY)
                ? context.get(AgentChatMemoryAdvisor.CHAT_MEMORY_AGENT_ID_KEY).toString() : "";
    }

    public static class Builder extends AbstractChatMemoryAdvisor.AbstractBuilder<ChatEnvironmentMemory> {

        protected Builder(ChatEnvironmentMemory chatEnvironmentMemory) {
            super(chatEnvironmentMemory);
        }

        public ChatExtendMemoryAdvisor build() {
            return new ChatExtendMemoryAdvisor(this.chatMemory, this.conversationId, this.chatMemoryRetrieveSize, this.order);
        }

    }
}